from langchain_community.document_loaders import DirectoryLoader
from major.models_manager import embedding_model
from langchain_chroma import Chroma

loader = DirectoryLoader("docs")
docs = loader.load()


vectorstore = Chroma.from_documents(
    docs,
    embedding=embedding_model.get_model(),
    persist_directory="./profile_chroma_db" # 注意是添加存储而非覆盖
)

results = vectorstore.similarity_search("夜神月", k=3)


for doc in results:
    print(f"内容: {doc.page_content}")
    print(f"来源: {doc.metadata['source']}")
    print("-" * 20)